import mlflow
import pandas as pd

import mlflow
import pandas as pd

def generate_recommendations_table(experiment_ids, aggregation_function="common_features", note="sizes_acts", group_type="sim"):
    all_rows = []

    for exp_id in experiment_ids:
        runs = mlflow.search_runs(
            experiment_ids=[exp_id],
            output_format="list"
        )
        for run in runs:
            if run.data.params.get("note") != note or run.data.params.get("SAE_fusion_strategy") != aggregation_function or run.data.params.get("group_type") != group_type:
                continue

            dataset = run.data.params.get("dataset", f"Exp-{exp_id}")
            dim = int(run.data.params.get("embedding_dim", 0))
            topk = int(run.data.params.get("top_k", 0))

            row_key = (dim, topk)
            metrics = {
                (dataset, "G/mean"): run.data.metrics.get("CommonItemsNDCG20/median"),
                (dataset, "U/mean"): run.data.metrics.get("NDCG20/mean"),
                (dataset, "U/min"): run.data.metrics.get("NDCG20/min"),
                (dataset, "Pop"): run.data.metrics.get("Popularity/mean"),
            }

            all_rows.append((row_key, metrics))

    # Build DataFrame from records
    records = {}
    for key, metrics in all_rows:
        if key not in records:
            records[key] = {}
        records[key].update(metrics)

    df = pd.DataFrame.from_dict(records, orient="index")
    df.index.names = ["Dimensions", "TopK"]

    # Sort and reindex columns by dataset then metric
    df = df.sort_index(axis=1, level=[0, 1]).sort_values(
        by=["Dimensions", "TopK"]
    )

    return df.reset_index()

def generate_sae_table(experiment_ids, note=None):
    all_rows = []

    for exp_id in experiment_ids:
        runs = mlflow.search_runs(
            experiment_ids=[exp_id],
            output_format="list"
        )
        for run in runs:
            if run.data.params.get("note") != note:
                continue

            dataset = run.data.params.get("dataset", f"Exp-{exp_id}")
            dim = int(run.data.params.get("embedding_dim", 0))
            base_factors = float(run.data.params.get("base_factors", 0.0))
            topk = int(run.data.params.get("top_k", 0))

            row_key = (dim, topk)
            metrics = {
                (dataset, "CS"): run.data.metrics.get("CosineSim/test"),
                (dataset, "Deg"): run.data.metrics.get("NDCG20_Degradation/test"),
                (dataset, "Deads"): run.data.metrics.get("DeadNeurons/test")
            }

            all_rows.append((row_key, metrics))

    # Build DataFrame from records
    records = {}
    for key, metrics in all_rows:
        if key not in records:
            records[key] = {}
        records[key].update(metrics)

    df = pd.DataFrame.from_dict(records, orient="index")
    df.index.names = ["Dimensions", "TopK"]

    # Sort and reindex columns by dataset then metric
    df = df.sort_index(axis=1, level=[0, 1]).sort_values(
        by=["Dimensions", "TopK"]
    )

    return df.reset_index()

import mlflow
import pandas as pd

def generate_common_features_table(experiment_ids, note="sizes_acts"):
    all_rows = []

    for exp_id in experiment_ids:
        runs = mlflow.search_runs(
            experiment_ids=[exp_id],
            filter_string=f"params.note = '{note}'",
            output_format="list"
        )

        for run in runs:
            params = run.data.params
            metrics = run.data.metrics
            
            if params.get("note") != note:
                continue
            

            dataset = params.get("dataset", f"Exp-{exp_id}")
            group_type = params.get("group_type", "unknown").capitalize()
            dim = int(params.get("embedding_dim", 0))
            topk = int(params.get("top_k", 0))
            value = metrics.get("common_features/mean", None)

            if value is None:
                continue

            row_key = (dim, topk)
            all_rows.append((row_key, (dataset, group_type), value))

    # Create dictionary for DataFrame
    records = {}
    for row_key, col_key, value in all_rows:
        if row_key not in records:
            records[row_key] = {}
        records[row_key][col_key] = value

    df = pd.DataFrame.from_dict(records, orient="index")
    df.index.names = ["Dimensions", "TopK"]

    # Reorder columns
    if not df.empty:
        datasets = ['LastFM1k',"MovieLens"]
        subcols = ["Sim", "Random", "Div"]
        col_order = [(d, s) for d in datasets for s in subcols]
        df = df.reindex(columns=pd.MultiIndex.from_tuples(col_order)).sort_values(
            by=["Dimensions", "TopK"]
        )

    return df.reset_index()

sae_experiments = ['657713966175362303', '852893065079987597']

table = generate_sae_table(sae_experiments, note="sizes_L2")
table.to_latex(
    "sae_table.tex",
    index=False,
    float_format="%.3f",
    column_format="lccc|ccc|ccc",
    escape=False,
    caption="SAE Table",
    label="tab:sae_table"
)
table

experiments = ['228719589483846826','962723054918039068']
generate_common_features_table(experiments, note="sizes_L2_with_acts")

experiment_ids = ['333391697323445885', '523100174176986081']

# nejdrive common_features
table = generate_recommendations_table(experiment_ids, aggregation_function="common_features", note="sizes_L2_with_acts")
table

table = generate_recommendations_table(experiment_ids, aggregation_function="average", note="sizes_L2_with_acts")
averaged_values = table.mean(axis=0, numeric_only=True)
table

generate_recommendations_table(experiment_ids, aggregation_function="common_features", note="sizes_L2_with_acts", group_type="random").round(3)

generate_recommendations_table(experiment_ids, aggregation_function="average", note="sizes_L2_with_acts", group_type="div")